/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2021-2021.
 * Description: support ksm_add feature, providing ksm.c
 *              to use some variables and functions.
 * Author: liwugui
 * Create: 2021-8-18
 */

#include <linux/errno.h>
#include <linux/ksm.h>
#include <linux/sched.h>
#include <linux/mm.h>
#include <linux/sched/mm.h>
#include <linux/task_struct_extend.h>
#include <asm-generic/mman-common.h>
#include <linux/proc_fs.h>
#include <linux/rtos_ksm_add.h>

#define ONE_PID 0x01                /* mark one pid */
#define ONE_PID_SCANNED 0x02        /* that task was scanned by for_each_process */
#define RECURSION_PID 0x04          /* mark recur pid */
#define RECURSION_PID_SCANNED 0x08  /* that task and its descendant were scanned */
#define ALL_PROCESS_SCANNED 0x10    /* that task was scanned by for_each_process */
#define EXCLUDE_PID 0x20            /* mark exclude pid */
#define RECUR_SCANNED 0x40          /* task_madvise_recur was scanned */
#define ALL_SCANNED  0x80           /* task_madvise_all was scanned */

#define ENABLE_ALL_PROCESS 1        /* it means ksm will merge all process */
#define ALL_PROCESS_BUFFER_LEN 2    /* all_process only accept single number */
#define DISABLE_KSM_ADD_DEBUG 0     /* close ksm_add debug printf */
#define ENABLE_KSM_ADD_DEBUG 1      /* open ksm_add debug printf */
#define KSM_ADD_DEBUG_BUFFER_LEN 2  /* ksm_add_debug only accept single number */
#define ALL_PROCESS 3               /* mark all process for scan_process, not present in ksm_add_flag */

#define DECIMAL 10                  /* the input PID can only be in decimal format. */
#define INPUT_JUDGE_LEN 4           /* the length of "add " and "del " */
#define INPUT_ADD 1                 /* the input parms is "add " */
#define INPUT_DEL 2                 /* the input parms is "del " */
#define BUFFER_LEN 15               /* maximum write length of the interface */

#define ksm_add_debug_pr(fmt, args...) \
	do { \
		if (ksm_add_debug) \
			pr_info("[KSM_ADD][DEBUG][%s][%d]: "fmt, __func__, __LINE__, ##args); \
	} while (0)

struct task_info {
	struct task_struct *task;
	struct list_head task_list;
};

static LIST_HEAD(task_list_head);
static DEFINE_MUTEX(ksm_scan_flag_mutex);

static bool all_process = DISABLE_ALL_PROCESS;
static bool ksm_add_debug = DISABLE_KSM_ADD_DEBUG;

/*
 * Checks the ksm_add_flag status. If the flag is set, return itself,
 * else, return 0.
 */
static inline unsigned int check_ksm_add_flag(struct task_struct *task, unsigned int flag)
{
	struct rtos_task_struct *rtos_task = task_to_rtos_task(task);

	return rtos_task->ksm_add_flag & flag;
}

static inline void set_ksm_add_flag(struct task_struct *task, unsigned int flag)
{
	struct rtos_task_struct *rtos_task = task_to_rtos_task(task);

	rtos_task->ksm_add_flag |= flag;
}

static inline void clear_ksm_add_flag(struct task_struct *task, unsigned int flag)
{
	struct rtos_task_struct *rtos_task = task_to_rtos_task(task);

	rtos_task->ksm_add_flag &= ~flag;
}

/* reset this flag of all process */
static void clear_all_ksm_add_flag(unsigned int flag)
{
	struct task_struct *task = NULL;

	rcu_read_lock();
	for_each_process(task) {
		clear_ksm_add_flag(task, flag);
	}
	rcu_read_unlock();
}

/*
 * Find mm through task and ksm_madvise all VMAs of the process.
 * task lock and get_task_struct will set out of the function.
 * exclude_pid is excluded in this function.
 */
static int task_madvise(struct task_struct *task, unsigned int advise)
{
	struct mm_struct *mm = NULL;
	struct vm_area_struct *vma = NULL;
	int err = 0;

	if ((advise == MADV_MERGEABLE) &&
	    check_ksm_add_flag(task, EXCLUDE_PID))
		return -EINVAL;

	/* get mm form task. if mm not NULL, get_task_mm will call mmget() */
	mm = get_task_mm(task);
	if (!mm)
		return -EINVAL;

	if (mmap_write_lock_killable(mm)) {
		err = -EINTR;
		goto lockfail;
	}

	vma = mm->mmap;
	/* madvise all vma, and if the mm_struct is the first madvise, it will call alloc_mm_slot */
	for (; vma; vma = vma->vm_next) {
		err = ksm_madvise(vma, vma->vm_start, vma->vm_end, (int)advise, &vma->vm_flags);
		if (err)
			goto madvfail;
	}
madvfail:
	mmap_write_unlock(mm);
lockfail:
	mmput(mm);
	return err;
}

static void free_task_list(void)
{
	struct task_info *task_info = NULL;
	struct task_info *task_info_tmp = NULL;

	list_for_each_entry_safe(task_info, task_info_tmp, &task_list_head, task_list) {
		clear_ksm_add_flag(task_info->task, RECUR_SCANNED);
		put_task_struct(task_info->task);
		list_del(&task_info->task_list);
		kfree(task_info);
	}
}

static int push_all_children(struct task_struct *task)
{
	struct task_info *task_info = NULL;
	struct task_struct *task_child = NULL;
	struct task_struct *madvise_task = NULL;
	struct list_head *list = NULL;
	int flag = 0;

	/* only all sons were added to list, flag would be true */
	while (!flag) {
		cond_resched();
		madvise_task = NULL;
		read_lock(&tasklist_lock);
		list_for_each(list, &task->children) {
			task_child = list_entry(list, struct task_struct, sibling);
			if (!check_ksm_add_flag(task_child, RECUR_SCANNED)) {
				set_ksm_add_flag(task_child, RECUR_SCANNED);
				madvise_task = get_task_struct(task_child);
				break;
			}
		}
		read_unlock(&tasklist_lock);

		if (!madvise_task) {
			flag++;
			break;
		}
		task_info = kmalloc(sizeof(struct task_info), GFP_KERNEL);
		if (!task_info) {
			put_task_struct(madvise_task);
			return -ENOMEM;
		}
		task_info->task = madvise_task;
		list_add_tail(&task_info->task_list, &task_list_head);
	}
	return 0;
}

/*
 * only recur_pid interface will call this function.
 * Madvise the current process and its descendants.
 */
static int task_madvise_recur(struct task_struct *task, unsigned int advise)
{
	struct task_info *task_info = NULL;
	int err;

	clear_all_ksm_add_flag(RECUR_SCANNED);
	task_info = kmalloc(sizeof(struct task_info), GFP_KERNEL);
	if (!task_info)
		return -ENOMEM;
	get_task_struct(task);
	set_ksm_add_flag(task, RECUR_SCANNED);
	task_info->task = task;
	list_add_tail(&task_info->task_list, &task_list_head);

	/* Using the queue bfs, iterate through the current process and its descendants. */
	while (!list_empty(&task_list_head)) {
		/* task_madvise queue head */
		task_info = list_entry(task_list_head.next, struct task_info, task_list);
		if (!((advise == MADV_UNMERGEABLE) &&
		    check_ksm_add_flag(task_info->task, ONE_PID)))
			task_madvise(task_info->task, advise);
		if (advise == MADV_MERGEABLE)
			set_ksm_add_flag(task_info->task, RECURSION_PID);
		else
			clear_ksm_add_flag(task_info->task, RECURSION_PID);

		/* create sons' task_info and add sons to the list tail. */
		err = push_all_children(task_info->task);
		if (err) {
			pr_info("[KSM_ADD][DEBUG][%s][%d]: push_all_children failed\n",
				__func__, __LINE__);
			free_task_list();
			return err;
		}

		/* pop queue head */
		task_info = list_entry(task_list_head.next, struct task_info, task_list);
		clear_ksm_add_flag(task_info->task, RECUR_SCANNED);
		put_task_struct(task_info->task);
		list_del(task_list_head.next);
		kfree(task_info);
	}
	return 0;
}

/*
 * only enable all_process and it will call this function task_madvise_all.
 * Select a process to madvise, in order to wake up ksm_scan_thread.
 */
static void task_madvise_all(void)
{
	struct task_struct *task = NULL;
	struct task_struct *madvise_task = NULL;
	int flag = 0;

	clear_all_ksm_add_flag(ALL_SCANNED);
	while (!flag) {
		cond_resched();
		madvise_task = NULL;
		rcu_read_lock();
		for_each_process(task) {
			if (!check_ksm_add_flag(task, ALL_SCANNED)) {
				set_ksm_add_flag(task, ALL_SCANNED);
				madvise_task = get_task_struct(task);
				break;
			}
		}
		rcu_read_unlock();

		/* if task_madvise is NULL, indicates it comes to the list tail */
		if (madvise_task) {
			if (all_process == ENABLE_ALL_PROCESS)
				task_madvise(madvise_task, MADV_MERGEABLE);
			if ((all_process == DISABLE_ALL_PROCESS) &&
			    !check_ksm_add_flag(madvise_task, ONE_PID) &&
			    !check_ksm_add_flag(madvise_task, RECURSION_PID))
				task_madvise(madvise_task, MADV_UNMERGEABLE);
			put_task_struct(madvise_task);
		} else {
			flag++;
			clear_all_ksm_add_flag(ALL_SCANNED);
		}
	}
}

/*
 * parms: ALL_PROCESS or ONE_PID or RECURSION_PID.
 * if ALL_PROCESS, traverse processes, select a process to register.
 * if ONE_PID, traverse the process linked list, and select
 * a proper process to madvise.
 * if RECURSION_PID, traverse the process linked list, and select
 * a process to madvise itself and its descendant.
 */
static void scan_process(unsigned int flag)
{
	struct task_struct *task = NULL;
	struct task_struct *madvise_task = NULL;
	unsigned int scanned;

	if (flag == ALL_PROCESS)
		scanned = ALL_PROCESS_SCANNED;
	else if (flag == ONE_PID)
		scanned = ONE_PID_SCANNED;
	else
		scanned = RECURSION_PID_SCANNED;

	/* find a unscanned process */
	rcu_read_lock();
	for_each_process(task) {
		if (((flag == ALL_PROCESS) ||
			check_ksm_add_flag(task, flag)) &&
			!check_ksm_add_flag(task, scanned)) {
			set_ksm_add_flag(task, scanned);
			madvise_task = get_task_struct(task);
			break;
		}
	}
	rcu_read_unlock();

	/* madvise process */
	if (madvise_task) {
		task_madvise(madvise_task, MADV_MERGEABLE);
		put_task_struct(madvise_task);
	} else {
		clear_all_ksm_add_flag(scanned);
	}
}

/*
 * Each time the ksm call this function to register
 * the newly generated process and vma.
 */
void do_ksm_add(void)
{
	mutex_lock(&ksm_scan_flag_mutex);
	if (all_process) {
		scan_process(ALL_PROCESS);
	} else {
		scan_process(ONE_PID);
		scan_process(RECURSION_PID);
	}
	mutex_unlock(&ksm_scan_flag_mutex);
}

/*
 * Check whether the input of interfaces one_pid or recur_pid is valid.
 */
static int check_parms(char *line, const char __user *buf, size_t len,
			unsigned long *pid_address, size_t arrlen)
{
	int err;

	if ((len < 1) || (len > BUFFER_LEN - 1) || (len > arrlen - 1)) {
		ksm_add_debug_pr("parms length illegal!\n");
		return -EINVAL;
	}

	memset(line, 0, BUFFER_LEN);
	if (copy_from_user(line, buf, len))
		return -EFAULT;

	if (!strncmp(line, "add ", INPUT_JUDGE_LEN)) {
		err = kstrtoul(line + INPUT_JUDGE_LEN, DECIMAL, pid_address);
		if (err || *pid_address > PID_MAX_LIMIT)
			return -EINVAL;
		return INPUT_ADD;
	}

	if (!strncmp(line, "del ", INPUT_JUDGE_LEN)) {
		err = kstrtoul(line + INPUT_JUDGE_LEN, DECIMAL, pid_address);
		if (err || *pid_address > PID_MAX_LIMIT)
			return -EINVAL;
		return INPUT_DEL;
	}

	return -EINVAL;
}

/*
 * Return to the user mode task based on the PID.
 */
static struct task_struct *find_get_user_task_by_vpid(pid_t pid)
{
	struct task_struct *task = NULL;

	/* if task is not NULL, find_get_task_by_vpid() will call get_task_struct() */
	task = find_get_task_by_vpid(pid);
	if (!task) {
		ksm_add_debug_pr("pid %d doesn't exist in the system\n", pid);
		return NULL;
	}

	if (!task->mm) {
		ksm_add_debug_pr("pid %d doesn't belong to the user mode\n", pid);
		put_task_struct(task);
		return NULL;
	}

	return task;
}

/*
 * parms: scan_flag value only are ONE_PID and RECURSION_PID
 * set ksm_add_flag and madvise MADV_MERGEABLE
 */
static int pid_add(pid_t pid, unsigned int scan_flag)
{
	int err = 0;
	struct task_struct *task = NULL;

	/* if task not NULL, it will call get_task_struct */
	task = find_get_user_task_by_vpid(pid);
	if (!task)
		return -EINVAL;

	mutex_lock(&ksm_scan_flag_mutex);
	set_ksm_add_flag(task, scan_flag);
	if (scan_flag == ONE_PID) {
		task_madvise(task, MADV_MERGEABLE);
	} else {
		err = task_madvise_recur(task, MADV_MERGEABLE);
		if (err)
			ksm_add_debug_pr("task_madvise_recur failed!\n");
	}
	mutex_unlock(&ksm_scan_flag_mutex);
	put_task_struct(task);

	return err;
}

/*
 * reset ksm_scan_flag and madvise MADV_UNMERGEABLE
 */
static int pid_del(pid_t pid, unsigned int flag)
{
	struct task_struct *task = NULL;

	task = find_get_task_by_vpid(pid);
	if (!task)
		return -EINVAL;

	mutex_lock(&ksm_scan_flag_mutex);

	if ((flag == ONE_PID) &&
	    (check_ksm_add_flag(task, ONE_PID) ||
	     check_ksm_add_flag(task, RECURSION_PID) ||
	     (all_process == ENABLE_ALL_PROCESS))) {
		set_ksm_add_flag(task, EXCLUDE_PID);
		clear_ksm_add_flag(task, ONE_PID);
		task_madvise(task, MADV_UNMERGEABLE);
	} else if ((flag == RECURSION_PID) &&
		   check_ksm_add_flag(task, RECURSION_PID)) {
		task_madvise_recur(task, MADV_UNMERGEABLE);
	} else {
		mutex_unlock(&ksm_scan_flag_mutex);
		put_task_struct(task);
		ksm_add_debug_pr("pid %d del failed!\n", pid);
		return -EINVAL;
	}
	mutex_unlock(&ksm_scan_flag_mutex);
	put_task_struct(task);

	return 0;
}

/*
 * The user input a pid and it will madvise this process.
 */
static ssize_t one_pid_write(struct file *file, const char __user *buf,
			     size_t len, loff_t *ppos)
{
	unsigned long pid_input;
	int ret;
	int err;
	char line[BUFFER_LEN];

	ret = check_parms(line, buf, len, &pid_input, BUFFER_LEN);
	if (ret < 0) {
		ksm_add_debug_pr("Invalid input format\n");
		return ret;
	}

	if (ret == INPUT_ADD)
		err = pid_add((pid_t)pid_input, ONE_PID);
	else
		err = pid_del((pid_t)pid_input, ONE_PID);
	if (err) {
		ksm_add_debug_pr("Invalid pid %d\n", (int)pid_input);
		return err;
	}

	ksm_add_debug_pr("%s pid %d success!\n",
			 (ret == INPUT_ADD) ? "add" : "del", (int)pid_input);
	return (ssize_t)len;
}

/*
 * The user input a pid and it will madvise the process and its descendants.
 */
static ssize_t recur_pid_write(struct file *file, const char __user *buf,
			     size_t len, loff_t *ppos)
{
	unsigned long pid_input;
	int ret;
	int err;
	char line[BUFFER_LEN];

	ret = check_parms(line, buf, len, &pid_input, BUFFER_LEN);
	if (ret < 0) {
		ksm_add_debug_pr("Invalid input format.\n");
		return ret;
	}

	if (ret == INPUT_ADD)
		err = pid_add((pid_t)pid_input, RECURSION_PID);
	else
		err = pid_del((pid_t)pid_input, RECURSION_PID);
	if (err) {
		ksm_add_debug_pr("Invalid pid %d\n", (int)pid_input);
		return err;
	}

	ksm_add_debug_pr("%s pid %d success in recur_pid!\n",
			 (ret == INPUT_ADD) ? "add" : "del", (int)pid_input);
	return (ssize_t)len;
}

/*
 * enable all_process and madvise all processes
 */
static ssize_t all_process_write(struct file *file, const char __user *buf,
			     size_t len, loff_t *ppos)
{
	int err;
	char line[BUFFER_LEN];
	unsigned long all_process_tmp;

	if (len != ALL_PROCESS_BUFFER_LEN) {
		ksm_add_debug_pr("The length of input parms is incorrect.\n");
		return -EINVAL;
	}

	memset(line, 0, BUFFER_LEN);
	if (copy_from_user(line, buf, len))
		return -EFAULT;

	err = kstrtoul(line, DECIMAL, &all_process_tmp);
	if (err || (all_process_tmp != ENABLE_ALL_PROCESS &&
	    all_process_tmp != DISABLE_ALL_PROCESS)) {
		ksm_add_debug_pr("Invalid input format\n");
		return -EINVAL;
	}

	if (all_process == (unsigned int)all_process_tmp) {
		ksm_add_debug_pr("Duplicate settings\n");
		return len;
	}
	all_process = (unsigned int)all_process_tmp;
	mutex_lock(&ksm_scan_flag_mutex);
	task_madvise_all();
	mutex_unlock(&ksm_scan_flag_mutex);

	ksm_add_debug_pr("%s all_process success!\n",
			 (all_process == ENABLE_ALL_PROCESS) ? "enable" : "disable");
	return (ssize_t)len;
}

static ssize_t ksm_add_debug_write(struct file *file, const char __user *buf,
			     size_t len, loff_t *ppos)
{
	int err;
	char line[BUFFER_LEN];
	unsigned long ksm_add_debug_tmp;

	if (len != KSM_ADD_DEBUG_BUFFER_LEN) {
		pr_info("[KSM_ADD][DEBUG][%s][%d]: The length of input parms is incorrect.\n",
			__func__, __LINE__);
		return -EINVAL;
	}

	memset(line, 0, BUFFER_LEN);
	if (copy_from_user(line, buf, len))
		return -EFAULT;

	err = kstrtoul(line, DECIMAL, &ksm_add_debug_tmp);
	if (err || (ksm_add_debug_tmp != ENABLE_KSM_ADD_DEBUG &&
	    ksm_add_debug_tmp != DISABLE_KSM_ADD_DEBUG)) {
		pr_info("[KSM_ADD][DEBUG][%s][%d]: Invalid input format\n",
			__func__, __LINE__);
		return -EINVAL;
	}

	ksm_add_debug = (unsigned int)ksm_add_debug_tmp;

	pr_info("[KSM_ADD][DEBUG][%s][%d]: %s ksm_add_debug success!\n",
		__func__, __LINE__,
		(ksm_add_debug == ENABLE_KSM_ADD_DEBUG) ? "enable" : "disable");
	return (ssize_t)len;
}

static int pid_show(struct seq_file *seq, unsigned int flag)
{
	struct task_struct *task = NULL;

	mutex_lock(&ksm_scan_flag_mutex);
	rcu_read_lock();
	for_each_process(task) {
		if (check_ksm_add_flag(task, flag))
			seq_printf(seq, "%d\n", task->pid);
	}
	rcu_read_unlock();
	mutex_unlock(&ksm_scan_flag_mutex);

	return 0;
}

static int one_pid_show(struct seq_file *seq, void *offset)
{
	pid_show(seq, ONE_PID);
	return 0;
}

static int one_pid_open(struct inode *inode, struct file *file)
{
	return single_open(file, one_pid_show, NULL);
}

static int recur_pid_show(struct seq_file *seq, void *offset)
{
	pid_show(seq, RECURSION_PID);
	return 0;
}

static int recur_pid_open(struct inode *inode, struct file *file)
{
	return single_open(file, recur_pid_show, NULL);
}

static int all_process_show(struct seq_file *seq, void *offset)
{
	seq_printf(seq, "%u\n", all_process);
	return 0;
}

static int all_process_open(struct inode *inode, struct file *file)
{
	return single_open(file, all_process_show, NULL);
}

static int ksm_add_debug_show(struct seq_file *seq, void *offset)
{
	seq_printf(seq, "%u\n", ksm_add_debug);
	return 0;
}

static int ksm_add_debug_open(struct inode *inode, struct file *file)
{
	return single_open(file, ksm_add_debug_show, NULL);
}

static const struct proc_ops one_pid_ops = {
	.proc_open = one_pid_open,
	.proc_read = seq_read,
	.proc_write = one_pid_write,
	.proc_lseek = seq_lseek,
	.proc_release = single_release,
};

static const struct proc_ops recur_pid_ops = {
	.proc_open = recur_pid_open,
	.proc_read = seq_read,
	.proc_write = recur_pid_write,
	.proc_lseek = seq_lseek,
	.proc_release = single_release,
};

static const struct proc_ops all_process_ops = {
	.proc_open = all_process_open,
	.proc_read = seq_read,
	.proc_write = all_process_write,
	.proc_lseek = seq_lseek,
	.proc_release = single_release,
};

static const struct proc_ops ksm_add_debug_ops = {
	.proc_open = ksm_add_debug_open,
	.proc_read = seq_read,
	.proc_write = ksm_add_debug_write,
	.proc_lseek = seq_lseek,
	.proc_release = single_release,
};

int ksm_add_create_dir(void)
{
	struct proc_dir_entry *proc_ksm_add = NULL;
	struct proc_dir_entry *proc_p = NULL;

	proc_ksm_add = proc_mkdir("ksm_add", NULL);
	if (!proc_ksm_add) {
		pr_err("Cannot create ksm_add proc dir.\n");
		goto err_ksm_add;
	}

	proc_p = proc_create("one_pid", 0640, proc_ksm_add, &one_pid_ops);
	if (!proc_p) {
		pr_err("Cannot create one_pid proc dir.\n");
		goto err_one_pid;
	}

	proc_p = proc_create("recur_pid", 0640, proc_ksm_add, &recur_pid_ops);
	if (!proc_p) {
		pr_err("Cannot create recur_pid proc dir.\n");
		goto err_recur_pid;
	}

	proc_p = proc_create("all_process", 0640, proc_ksm_add, &all_process_ops);
	if (!proc_p) {
		pr_err("Cannot create all_process proc dir.\n");
		goto err_all_process;
	}

	proc_p = proc_create("ksm_add_debug", 0640, proc_ksm_add, &ksm_add_debug_ops);
	if (!proc_p) {
		pr_err("Cannot create ksm_add_debug proc dir.\n");
		goto err_ksm_add_debug;
	}
	return 0;
err_ksm_add_debug:
	remove_proc_entry("all_process", proc_ksm_add);
err_all_process:
	remove_proc_entry("recur_pid", proc_ksm_add);
err_recur_pid:
	remove_proc_entry("one_pid", proc_ksm_add);
err_one_pid:
	remove_proc_entry("ksm_add", NULL);
err_ksm_add:
	return -ENOMEM;
}

/*
 * when process fork, if father process belongs to RECURSION_PID,
 * son process will be setted to RECURSION_PID.
 */
void ksm_add_flag_inher(struct task_struct *task_parent, struct task_struct *task_child)
{
	struct rtos_task_struct *rtos_task_ksm_add_parent = NULL;
	struct rtos_task_struct *rtos_task_ksm_add_child = NULL;

	rtos_task_ksm_add_parent = task_to_rtos_task(task_parent);
	rtos_task_ksm_add_child = task_to_rtos_task(task_child);
	rtos_task_ksm_add_child->ksm_add_flag = rtos_task_ksm_add_parent->ksm_add_flag & RECURSION_PID;
}
